# gc();rm(list = ls());
gc()

# library(dplyr)
# library(ggplot2)
# library(tidyr)

load("results/ep/full.Rda")

dir.create("plots/ep/", recursive = TRUE)

supp.labs <- c("LR Dimension", "EU Dimension")
names(supp.labs) <- c("1", "2")
coef.labs <- c("Group Indicator", "Ideal Point (t-1)")
names(coef.labs) <- c("2", "3")

# nonseparability plot
ggplot(ns_est_full,
       aes(x = session, y = mean, ymin = q5, ymax = q95)) + geom_point(size = 3) + geom_linerange() + theme_bw() +
  xlab("Session") + ylab("Non-separability") + geom_hline(aes(yintercept = 0), linetype = 2, size = .5) +
  geom_smooth(aes(linetype = parameter), color = "black", size = .5) +
  theme(legend.position = "bottom", legend.title = element_blank()) + scale_linetype_discrete(labels = c("Non-separable")) + scale_shape_discrete(labels = c("Non-separable"))
ggsave(filename = "plots/ep/nsep_e.pdf", width = 4, height = 4)

# salience plot 
sal_est_full %>% select(mean, q5, q95, model, id, session) %>% 
  ggplot(aes(x = session, y = mean, ymin = q5, ymax = q95, shape = model)) + geom_point(size = 3, position = position_dodge(width = .3)) + geom_linerange(position = position_dodge(width = .3)) + theme_bw() +
  xlab("Session") + ylab("Salience") + facet_wrap(~id, ncol = 1, labeller = labeller(id = supp.labs)) + geom_smooth(aes(linetype = model), position = position_dodge(width = .3), color = "black", size = .5) + 
  theme(legend.position = "bottom", legend.title = element_blank()) + scale_linetype_discrete(labels = c("Non-separable", "Separable")) + scale_shape_discrete(labels = c("Non-separable", "Separable"))
ggsave(filename = "plots/ep/sal_e.pdf", width = 4, height = 6)

# correlation plot
cor_est_full %>% select(mean, q5, q95, model, session) %>% 
  ggplot(aes(x = session, y = mean, ymin = q5, ymax = q95, shape = model)) + geom_point(size = 3, position = position_dodge(width = .3)) + geom_linerange(position = position_dodge(width = .3)) + theme_bw() +
  xlab("Session") + ylab("Correlation across Dimensions") + geom_smooth(aes(linetype = model), position = position_dodge(width = .3), color = "black", size = .5) + 
  geom_hline(aes(yintercept = 0), linetype = 2) + 
  theme(legend.position = "bottom", legend.title = element_blank()) + scale_linetype_discrete(labels = c("Non-separable", "Separable")) + scale_shape_discrete(labels = c("Non-separable", "Separable"))
ggsave(filename = "plots/ep/cor_e.pdf", width = 4, height = 4)

# correlation plot
old_est_full %>% select(mean, q5, q95, model, id, session) %>% filter(session > 1) %>% 
  ggplot(aes(x = session, y = mean, ymin = q5, ymax = q95, shape = model)) + geom_point(size = 3, position = position_dodge(width = .3)) + geom_linerange(position = position_dodge(width = .3)) + theme_bw() +
  xlab("Session") + ylab("Correlation over Time") + facet_wrap(~id, ncol = 1, labeller = labeller(id = supp.labs)) + geom_smooth(aes(linetype = model), position = position_dodge(width = .3), color = "black", size = .5) + 
  theme(legend.position = "bottom", legend.title = element_blank()) + scale_linetype_discrete(labels = c("Non-separable", "Separable")) + scale_shape_discrete(labels = c("Non-separable", "Separable"))
ggsave(filename = "plots/ep/corold_e.pdf", width = 4, height = 6)

# theta plot
id_est_full %>% select(mean, model, id, dim, session) %>% 
  pivot_wider(names_from = dim, values_from = mean) %>% 
  rename(d1 = `1`, d2 = `2`) %>% 
  filter(model == "sep") %>% 
  ggplot(aes(d1, d2))+geom_point(alpha = .25) + xlab("LR Dimension") + ylab("EU Dimension") + theme_bw() + facet_wrap(~session, ncol = 5) + geom_smooth(method = "lm", color = "black", alpha = .2, size = .5)
ggsave(filename = "plots/ep/ideal_sep_e.pdf", width = 10, height = 4)
id_est_full %>% select(mean, model, id, dim, session) %>% 
  pivot_wider(names_from = dim, values_from = mean) %>% 
  rename(d1 = `1`, d2 = `2`) %>% 
  filter(model != "sep") %>% 
  ggplot(aes(d1, d2))+geom_point(alpha = .25) + xlab("LR Dimension") + ylab("EU Dimension") + theme_bw() + facet_wrap(~session, ncol = 5) + geom_smooth(method = "lm", color = "black", alpha = .2, size = .5)
ggsave(filename = "plots/ep/ideal_nsep_e.pdf", width = 10, height = 4)

# lm plots
gamma_est_full %>% 
  filter(parameter == "r_sq") %>% 
  select(mean, q5, q95, model, dim, session) %>% 
  ggplot(aes(x = session, y = mean, ymin = q5, ymax = q95, shape = model)) + geom_point(position = position_dodge(width = .3)) + geom_linerange(position = position_dodge(width = .3)) + theme_bw() +
  xlab("Session") + ylab("R^2") + facet_wrap(~dim, ncol = 1, labeller = labeller(dim = supp.labs)) + geom_smooth(aes(linetype = model), position = position_dodge(width = .3), color = "black", size = .5) + 
  theme(legend.position = "bottom", legend.title = element_blank()) + scale_linetype_discrete(labels = c("Non-separable", "Separable")) + scale_shape_discrete(labels = c("Non-separable", "Separable"))
ggsave(filename = "plots/ep/rsq_e.pdf", width = 4, height = 6)

gamma_est_full %>% 
  filter(parameter != "r_sq") %>% 
  mutate(dim = rep(c(1,2), length.out = 56),
         coef = c(rep(rep(c(1,2), each = 2), length.out = 8),
                  rep(rep(c(1,2,3), each = 2), length.out = 48))) %>% 
  filter(coef > 1) %>% 
  select(mean, q5, q95, model, dim, session, coef) %>% 
  ggplot(aes(x = session, y = mean, ymin = q5, ymax = q95, shape = model)) + geom_point(position = position_dodge(width = .3)) + geom_linerange(position = position_dodge(width = .3)) + theme_bw() +
  xlab("Session") + ylab("Coefficient") + facet_grid(dim~coef, labeller = labeller(dim = supp.labs, coef = coef.labs)) + 
  geom_smooth(aes(linetype = model), position = position_dodge(width = .3), color = "black", size = .5) + 
  theme(legend.position = "bottom", legend.title = element_blank()) + scale_linetype_discrete(labels = c("Non-separable", "Separable")) + scale_shape_discrete(labels = c("Non-separable", "Separable"))
ggsave(filename = "plots/ep/coef_e.pdf", width = 6, height = 6)

names(supp.labs) <- c("d1", "d2")
id_est_full %>% select(mean, model, id, dim, session) %>% 
  pivot_wider(names_from = dim, values_from = mean) %>% 
  rename(d1 = `1`, d2 = `2`) %>% 
  pivot_wider(names_from = model, values_from = c("d1", "d2")) %>% 
  group_by(session) %>% 
  summarize(d1 = summary(lm(d1_nonsep ~ d1_sep))$r.squared,
            d2 = summary(lm(d2_nonsep ~ d2_sep))$r.squared) %>% 
  ungroup() %>% 
  pivot_longer(cols = c("d1", "d2")) %>% 
  ggplot(aes(x = session, y = value)) + geom_point(size = 3) + facet_wrap(~name, ncol = 1, labeller = labeller(name = supp.labs)) + geom_smooth(size = .5, color = "black") + theme_bw() +
  xlab("Session") + ylab("Correlation across Specifications")
ggsave(filename = "plots/ep/cor_spec_e.pdf", width = 4, height = 6)

# loo plots
loo_est_full %>% 
  filter(elpd_diff != 0) %>% 
  ggplot(aes(x = session, y = elpd_diff, ymin = elpd_diff - 2 * se_diff, ymax = elpd_diff + 2 * se_diff, shape = model, linetype = model)) + 
  geom_linerange(position = position_dodge(width = .3)) + 
  geom_point(size = 3, position = position_dodge(width = .3)) + 
  theme_bw() + xlab("Non-separability") + ylab("LOO Comparison") + 
  theme(legend.position = "bottom", legend.title = element_blank()) + 
  scale_linetype_discrete(labels = c("Non-separable", "Separable")) + scale_shape_discrete(labels = c("Non-separable", "Separable"))
ggsave(filename = "plots/ep/loo.pdf", width = 4, height = 4)

id_est_full %>% 
  select(mean, model, id, dim, session) %>% 
  group_by(model, dim, session) %>% 
  mutate(mean = mean - mean(mean)) %>%
  pivot_wider(id_cols = c("model", "session", "id"), names_from = "dim", values_from = "mean") %>% 
  rename(d1 = `1`, d2 = `2`) %>% 
  pivot_wider(names_from = "model", values_from = c("d1", "d2")) %>% 
  mutate(dist = sqrt((d1_sep - d1_nonsep)^2 + (d2_sep - d2_nonsep)^2)) %>% 
  group_by(session) %>% 
  summarize(dist = mean(dist, na.rm = T)) %>% 
  merge(ns_est_full %>% select(session, mean)) %>% 
  ggplot(aes(mean, dist, label = session)) + geom_point() + theme_bw() + geom_label() + xlab("Non-separability") + ylab("Mean Distance")
ggsave(filename = "plots/ep/dist.pdf", width = 4, height = 4)

id_est_full %>% 
  select(mean, model, id, dim, session) %>% 
  group_by(model, dim, session) %>% 
  mutate(mean = mean - mean(mean)) %>%
  pivot_wider(id_cols = c("model", "session", "id"), names_from = "dim", values_from = "mean") %>% 
  rename(d1 = `1`, d2 = `2`) %>% 
  pivot_wider(names_from = "model", values_from = c("d1", "d2")) %>% 
  mutate(dist = sqrt((d1_sep - d1_nonsep)^2 + (d2_sep - d2_nonsep)^2)) %>% 
  ggplot(aes(d1_sep, d2_sep, xend = d1_nonsep, yend = d2_nonsep)) + geom_segment(arrow = arrow(length = unit(0.03, "npc"))) + theme_bw() + facet_wrap(~session, ncol = 5) + 
  xlab("LR Dimension") + ylab("EU Dimension")
ggsave(filename = "plots/ep/dist2.pdf", width = 10, height = 4)

# prediction and complexity
prcomp_est_full %>% 
  filter(parameter == "p_share") %>% 
  select(session, model, mean) %>% 
  pivot_wider(values_from = mean, names_from = model) %>% 
  ggplot(aes(x = nonsep, y = sep)) + 
  geom_point(alpha = .25) + 
  geom_abline(aes(slope = 1, intercept = 0)) +
  theme_bw() + xlab("Non-separable Specification") + ylab("Separable Specification") + 
  theme(legend.position = "bottom", legend.title = element_blank())
ggsave(filename = "plots/ep/pred.pdf", width = 3.5, height = 3.5)

prcomp_est_full %>% 
  filter(parameter == "m_prop_compl") %>% 
  select(session, model, mean) %>% 
  pivot_wider(values_from = mean, names_from = model) %>% 
  ggplot(aes(x = nonsep, y = sep)) + 
  geom_point(alpha = .25) + 
  theme_bw() + xlab("Non-separable Specification") + ylab("Separable Specification") + 
  theme(legend.position = "bottom", legend.title = element_blank()) + 
  geom_abline(aes(slope = 1, intercept = 0))
ggsave(filename = "plots/ep/compl.pdf", width = 3.5, height = 3.5)
